#!/usr/bin/env python

import os
import sys
import cPickle
import numpy as np
from datetime import datetime, timedelta

from INFILE import siteid,maindir,webdir
sys.path.append('%s/data_parsers' % maindir)


export_flag = True
ndays = 10   # How many days we want to plot back


def main():
    print "###################################"
    print "  Beginning to plot precip errors"
    print "###################################"
    truth, modd = get_data()
    plot_data(truth,modd)


def plot_data(truth, mod):
    import matplotlib
    if export_flag:
        matplotlib.use('agg')
    import matplotlib.pyplot as plt

    # loop through the models we care about
    for model in ['12Z GFS','12Z NAM','HPC','NWS','PW HRMOS','USL12Z','USL22Z']:
        try:
            datelist = mod[model].keys()
        except:
            continue
        datelist.sort()
        # Only the most recent ndays
        datelist = datelist[-(ndays+1):]

        # Now the model values rounded to nearest hundredth
        medianlist = [round(mod[model][d],2) for d in datelist] 


        # Reset the datelist to be just the days
        new_datelist = [datetime(d.year,d.month,d.day) for d in datelist]
        datelist = new_datelist
        #print model
        #print len(datelist), "DATELIST"
        #print len(medianlist), "MODELLIST"
        
        # Now get the truth
        truthlist = [truth[d]['precip'] for d in datelist[:-1]]
        # The last day doesn't exist yet
        truthlist.append(np.nan)
        #print len(truthlist), "TRUTHLIST"

        # IF the model is PW HRMOS, also get the DC HRMOS
        if model == 'PW HRMOS':
            medianlist_2 = [round(mod['DC HRMOS'][d],2) for d in datelist]



        # Make a plot
        plt.figure(figsize=(12,11))
        plt.subplot(2,1,1)
        truplot = plt.plot(datelist,truthlist,'k-',linewidth=4,gid='truth')
        plt.hold(True)
        medplot = plt.plot(datelist,medianlist,'b-',linewidth=3,gid=model)

        #point out tomorrow's forecast
        tom_hpc = medianlist[-1]
        plt.annotate('%1.2f'%(tom_hpc), xy=(matplotlib.dates.date2num(datelist[-1]),tom_hpc),xycoords='data', \
            textcoords='offset points', xytext=(-10,25),size='medium',va="center",ha="left",bbox=dict(boxstyle="round4", fc="w", edgecolor='b'),\
            arrowprops=dict(arrowstyle="->",edgecolor='b'))

        # If we're doing PW HRMOS, then also do DC HRMOS
        if model == 'PW HRMOS':
            medplot_2 = plt.plot(datelist,medianlist_2,'g-',linewidth=3,gid='DC HRMOS')
            tom_hpc = medianlist_2[-1]
            plt.annotate('%1.2f'%(tom_hpc), xy=(matplotlib.dates.date2num(datelist[-1]),tom_hpc),xycoords='data', \
                textcoords='offset points', xytext=(-10,25),size='medium',va="center",ha="left",bbox=dict(boxstyle="round4", fc="w", edgecolor='g'),\
                arrowprops=dict(arrowstyle="->",edgecolor='g'))
            leg = plt.legend([medplot[0],medplot_2[0],truplot[0]],[model,'DC HRMOS','Actual'],loc=2)
            plt.title('Past %d-day forecast precip comparison at %s (12Z HRMOS)' % (ndays,siteid.upper()))
        else:
            leg = plt.legend([medplot[0],truplot[0]],[model,'Actual'],loc=2)
            plt.title('Past %d-day forecast precip comparison at %s (%s)' % (ndays,siteid.upper(),model))

        ymin,ymax = plt.ylim()
        plt.ylim(ymin = 0.0)
        ltexts = leg.get_texts()
        plt.setp(ltexts,fontsize='small')
        ax = plt.gca()
        ax.xaxis.set_major_formatter(matplotlib.dates.DateFormatter('%b %d'))
        #ax.set_yticks(range(int(ymin)-1,int(ymax)+1))
        plt.grid()
        plt.ylabel('Precipitation (in.)')

        # Now the error bars
        plt.subplot(2,1,2)
        modelerror = np.subtract(medianlist,truthlist)
        # Find mean error
        masked_model = np.ma.masked_array(modelerror,np.isnan(modelerror))
        meanf_model_error = np.mean(masked_model) 
        # Convert to string
        if meanf_model_error >= 0.0:
            mean_model_error = '+%1.2f' % meanf_model_error
        else:
            mean_model_error = '%1.2f' % meanf_model_error

        if model == 'PW HRMOS':
            modelerror_2 = np.subtract(medianlist_2,truthlist)
            # Find mean error
            masked_model = np.ma.masked_array(modelerror_2,np.isnan(modelerror_2))
            meanf_model_error = np.mean(masked_model) 
            # Convert to string
            if meanf_model_error >= 0.0:
                mean_model_error_2 = '+%1.2f' % meanf_model_error
            else:
                mean_model_error_2 = '%1.2f' % meanf_model_error

 

        # Now some fiddling to make the bar chart work
        try:
            dxval = abs(matplotlib.dates.date2num(datelist[1]) - matplotlib.dates.date2num(datelist[0]))
        except:
            # We don't have enough dates yet for this to work
            print "MODEL failed:", model
            continue
        datelistmod = [d - timedelta(hours=12) for d in datelist]   # This will offset the left side of the bar
        model_error_bar = plt.bar(datelistmod,modelerror,color='b',align='edge',bottom=0,width=dxval,gid='model_error')
        if model == 'PW HRMOS':
            dxval = dxval/2.
            datelistmod_2 = [d-timedelta(hours=6) for d in datelist]
            model_error_bar_2 = plt.bar(datelistmod_2,modelerror_2,color='g',align='edge',bottom=0,width=dxval,alpha=0.5,gid='model_error_2')

        plt.xlim(xmax=matplotlib.dates.date2num(datelistmod[-1]+timedelta(hours=12)))
        plt.xlim(xmin=matplotlib.dates.date2num(datelistmod[0]+timedelta(hours=12)))
        plt.hold(True)
        # Zero-line
        zeroline = plt.axhline(y=0,color='k',linestyle='--',linewidth='4')
 
        # Display the mean text
        if model == 'PW HRMOS':
            txtb = 'Mean %s error: %s | Mean DC HRMOS error: %s' % (model,mean_model_error,mean_model_error_2)
            leg = plt.legend([zeroline,model_error_bar[0],model_error_bar_2[0]],['Zero',model,'DC HRMOS'])
        else:
            txtb = 'Mean %s error: %s' % (model,mean_model_error)
            leg = plt.legend([zeroline,model_error_bar[0]],['Zero',model])

        ax = plt.gca()
        plt.text(0.5,0.90,txtb,fontsize=14,horizontalalignment='center',transform=ax.transAxes,bbox=dict(facecolor='gray',alpha=0.1))

        plt.title('Magnitude of error in %s forecast' % model)
        ax.xaxis.set_major_formatter(matplotlib.dates.DateFormatter('%b %d'))
        plt.grid()
        plt.ylabel('Precip error (in.)')
        ltexts = leg.get_texts()
        plt.setp(ltexts,fontsize='small')

   

        if export_flag:
            if model == 'PW HRMOS':
                savelen=5
            elif model.startswith('USL'):
                savelen = len(model)
            else:
                savelen=3
            plt.savefig('%s_%s_precip_%dday.png' % (siteid.upper(),model[-savelen:],ndays), bbox_inches='tight')
            os.system('mv %s_%s_precip_%dday.png %s' % (siteid.upper(),model[-savelen:],ndays,webdir))
        else:
            plt.show()

        



def get_data():
    # Load the precip
    precipd = cPickle.load(open('./site_data/%s_precip_fcsts.pickle' % (siteid.upper()),'r'))

    # And load the HPC
    try:
        hpcd = cPickle.load(open('./hpc_pqpf/%s_hpc_archive.pickle' % siteid.upper(), 'r'))
        # For the HPC need to quickly convert the percentiles to median
        new_hpc = {}
        for date in hpcd.keys():
            new_hpc[date] = np.median(hpcd[date])

        # Add the HPC to the master dictionary
        precipd['HPC'] = new_hpc
    except:
        print "Unable to load HPC precip"

    # Now we need the PW Precip and DC Precip
    try:
        hrmosd = cPickle.load(open('./hrmos_qpf/%s_hrmos_dict.pickle' % siteid.upper(),'r'))
        # Add these to the dictionary, but change the names
        named = {'PW Precip' : 'PW HRMOS',
                 'DC Precip' : 'DC HRMOS'}
        for mod in hrmosd.keys():
            precipd[named[mod]] = hrmosd[mod]
    except:
        print "Unable to load HR MOS precip"

    # Also try the NWS data
    try:
        f = open('./site_data/%s_NWS_fcst.pickle' % siteid.upper(),'r')
        nws = cPickle.load(f)
        f.close()
        new_nws = {}
        for date in nws.keys():
            new_nws[date] = nws[date]['precip']
        # Add the NWS to the directory
        precipd['NWS'] = new_nws
    except:
        print "Unable to load NWS"

    # Finally the USL data
    try:
        f = open('./site_data/%s_USL12Z_fcst.pickle' % siteid.upper(), 'r')
        usl12Z = cPickle.load(f)
        f.close()
        f = open('./site_data/%s_USL22Z_fcst.pickle' % siteid.upper(), 'r')
        usl22Z = cPickle.load(f)
        f.close()
        new_12Z = {}
        new_22Z = {}
        for date in usl12Z.keys():
            new_12Z[date] = usl12Z[date]['precip']
        for date in usl22Z.keys():
            new_22Z[date] = usl22Z[date]['precip']
        # Add these to the master directory
        precipd['USL12Z'] = new_12Z
        precipd['USL22Z'] = new_22Z
    except:
        print "Unable to load USL data"

        
    
    # Now get the truth
    from encode_truth import encode_truth
    truthd = encode_truth(siteid.upper())

    return truthd, precipd

if __name__ == '__main__':
    main()

